package com.bleeth.flow.step.row;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import com.bleeth.flow.core.common.APlugin;
import com.bleeth.flow.core.common.PluginAnnotation;
import com.bleeth.flow.core.common.PluginTypeEnum;
import com.jd.platform.async.worker.WorkResult;
import com.jd.platform.async.wrapper.WorkerWrapper;
import lombok.Data;
import org.apache.spark.sql.Column;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.functions;

import java.io.Serializable;
import java.util.List;
import java.util.Map;

/**
 * @author ：Bleeth
 * @date ：2021-08-06
 * @description：取样组件
 *
 */
@Data
@PluginAnnotation(name = "取样组件",
        type = PluginTypeEnum.ROW,
        description = "",
        id = "SamplePlugin")
public class SamplePlugin extends APlugin implements Serializable {

    private static final long serialVersionUID = -5281128600781056098L;


    private List<SortParameter> sortList;


    @Override
    public void result(boolean success, Dataset<Row> param, WorkResult<Dataset<Row>> workResult) {

    }

    @Override
    public Dataset<Row> action(Dataset<Row> param, Map<String, WorkerWrapper> allWrappers) {
        String fromPluginName = fromList.get(0);
        WorkerWrapper fromWrapper = allWrappers.get(fromPluginName);
        Dataset<Row> ds = (Dataset<Row>) fromWrapper.getWorkResult().getResult();

        List<Column> sortColList = CollUtil.newArrayList();
        for (int i = 0; i < sortList.size(); i++) {
            SortParameter sortParameter = sortList.get(i);
            String fieldName = sortParameter.getFieldName();
            String sortType = sortParameter.getSortType();
            if (StrUtil.equalsIgnoreCase(sortType, "ASC")) {
                Column column = functions.col(fieldName).asc();
                sortColList.add(column);
            }
            if (StrUtil.equalsIgnoreCase(sortType, "DESC")) {
                Column column = functions.col(fieldName).desc();
                sortColList.add(column);
            }
        }
        ds = ds.sort(ArrayUtil.toArray(sortColList, Column.class));
        return ds;
    }


    @Data
    public static class SortParameter implements Serializable {

        private static final long serialVersionUID = 1028796001507450553L;

        private String fieldName;

        //排序类型（ASC\DESC）
        private String sortType;

    }


}
